import torch
from torch import nn
import numpy as np
from auxmodels.core.model import NNLoss
from auxmodels.utils.calc import xywh2xyxy

from ..utils.cal_funcs import decode_bbox, sigmoid_np


class CenterLoss(NNLoss):
    """损失函数"""

    def __init__(self, sub_loss_names, sub_loss_weights=None, stride=4, scale=64):
        super(CenterLoss, self).__init__(sub_loss_names, sub_loss_weights)
        self.conf_loss = HeatmapLoss()
        self.loc_loss = DIOULoss()
        self.cls_loss = nn.CrossEntropyLoss()
        self.stride = stride
        self.scale = scale

    def calculate_loss(self, outputs, targets):
        """
        targets:
            0:4 - bbox
            4:4+num_cls - heatmap
            4+num_cls - ignore
            4+num_cls+1 - rotate
        """
        out_bbox, out_hm, out_pool = outputs

        t_bbox = targets[..., :4]
        t_conf = targets[..., 4:-1]
        t_ignore = targets[..., -1]

        mask_target = (t_conf > 0.9999).sum(dim=-1).bool()
        mask_bbox = ~torch.isnan(t_bbox[..., 0])

        mask_ignore = ~torch.isnan(t_ignore)

        t_bbox = t_bbox[mask_bbox]
        b, row, col = torch.where(mask_bbox)

        # 置信度损失
        loss_conf = self.conf_loss(out_hm.permute(0, 2, 3, 1).contiguous(), t_conf)

        bbox = decode_bbox(
            torch.tanh(out_bbox.permute(0, 2, 3, 1).contiguous()[mask_bbox]),
            row,
            col,
            self.stride,
            self.stride * self.scale,
            torch.stack,
        )
        # 定位损失
        loss_loc = self.loc_loss(bbox, t_bbox)

        return (
            loss_conf * self.weights[0],
            loss_loc * self.weights[1],
        )


def focal_loss(
    output, target, t_ignore, classes_weights_pos, classes_weights_neg, eps=1e-5, gamma=1
):
    """
    :param output:
    :param target:
    :param t_bbox:
    :param t_ignore:
    :param eps:
    :param gamma:  值越大，惩罚越小
    :return:
    """

    center_pos_mask = target == 1.0

    signed = True
    gap = 5
    range_bits = 8 - signed - gap
    high = sigmoid_np(2**range_bits)
    low = 1 - high
    output = torch.clip(output, min=low, max=high)
    target = torch.clip(target, min=low, max=high)

    # 此处 应该取
    # center_neg_mask = torch.isnan(t_cls).unsqueeze(3).repeat(1, 1, 1, 6)  # & neg_mask
    # print(center_pos_mask.shape[-1])
    center_ignore_mask = (
        ~torch.isnan(t_ignore).unsqueeze(3).repeat(1, 1, 1, center_pos_mask.shape[-1])
    )
    center_neg_mask = ~torch.logical_or(center_pos_mask, center_ignore_mask)
    neg_weights = torch.pow(high - target, 2) * center_neg_mask.float() * 4.0  # lt 为小于

    pos_loss = torch.log(output + low) * torch.pow(high - output, gamma) * center_pos_mask.float()

    neg_loss = torch.log(1 + low - output) * torch.pow(output - low, gamma) * neg_weights

    if classes_weights_pos:
        for cls, weight in enumerate(classes_weights_pos):
            pos_loss[:, :, :, cls] = pos_loss[:, :, :, cls] * weight
    if classes_weights_neg:
        for cls, weight in enumerate(classes_weights_neg):
            neg_loss[:, :, :, cls] = neg_loss[:, :, :, cls] * weight

    num_pos, pos_loss, neg_loss = center_pos_mask.sum(), pos_loss.sum(), neg_loss.sum()
    loss = 0
    if num_pos == 0:
        loss = loss - neg_loss
    else:
        loss = loss - (pos_loss + neg_loss) / num_pos
    return loss


class HeatmapLoss(nn.Module):
    def __init__(self, alpha=[1, 1], gamma=4):
        super(HeatmapLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, pred, gt):
        eps = 1e-7
        pred = torch.sigmoid(pred)
        pos_inds = gt.eq(1).float()
        neg_inds = gt.lt(1).float()
        neg_weights = torch.pow(1 - gt, self.gamma)
        loss = torch.tensor(0.0, requires_grad=True)
        pos_loss = torch.log(pred + eps) * torch.pow(1 - pred, 2) * pos_inds
        neg_loss = torch.log(1 - pred + eps) * torch.pow(pred, 2) * neg_weights * neg_inds
        num_pos = pos_inds.sum()
        pos_loss = self.alpha[0] * pos_loss.sum()
        neg_loss = self.alpha[1] * neg_loss.sum()
        if num_pos == 0:
            loss = loss - neg_loss
        else:
            loss = loss - (pos_loss + neg_loss) / num_pos
            # print(pos_loss, neg_loss, num_pos)
        return loss


class DIOULoss(torch.nn.Module):
    def __init__(self, eps=1e-7, reduction="mean"):
        super(DIOULoss, self).__init__()
        self.eps = eps
        self.reduction = reduction

    def forward(self, preds, bbox):
        cx1 = preds[:, 0]
        cy1 = preds[:, 1]
        w1 = preds[:, 2]
        h1 = preds[:, 3]

        cx2 = bbox[:, 0]
        cy2 = bbox[:, 1]
        w2 = bbox[:, 2]
        h2 = bbox[:, 3]

        preds = xywh2xyxy(preds)
        bbox = xywh2xyxy(bbox)

        ix1 = torch.max(preds[:, 0], bbox[:, 0])
        iy1 = torch.max(preds[:, 1], bbox[:, 1])
        ix2 = torch.min(preds[:, 2], bbox[:, 2])
        iy2 = torch.min(preds[:, 3], bbox[:, 3])

        iw = (ix2 - ix1).clamp(min=0.0)
        ih = (iy2 - iy1).clamp(min=0.0)

        # overlaps
        inters = iw * ih

        # union
        uni = w1 * h1 + w2 * h2 - inters

        # iou
        iou = inters / (uni + self.eps)

        # inter_diag
        inter_diag = (cx1 - cx2) ** 2 + (cy1 - cy2) ** 2

        # outer_diag
        ox1 = torch.min(preds[:, 0], bbox[:, 0])
        oy1 = torch.min(preds[:, 1], bbox[:, 1])
        ox2 = torch.max(preds[:, 2], bbox[:, 2])
        oy2 = torch.max(preds[:, 3], bbox[:, 3])

        outer_diag = (ox1 - ox2) ** 2 + (oy1 - oy2) ** 2

        diou = iou - inter_diag / outer_diag

        diou = torch.clamp(diou, min=-1.0, max=1.0)

        diou_loss = 1.0 - diou

        if self.reduction == "mean":
            loss = torch.mean(diou_loss)
        elif self.reduction == "sum":
            loss = torch.sum(diou_loss)
        else:
            raise NotImplementedError
        return loss
